"""
    GPSRO

Observation operators for GPS Radio Occultation observations.

GPS Radio Occultation (GPS-RO) provides vertical profiles of:
- Refractivity (N): Related to atmospheric density
- Bending angle (α): Ray bending through atmosphere

# Physical Principles

## Refractivity
Atmospheric refractivity is computed from temperature, pressure, and humidity:

    N = 77.6·(p/T) + 3.73×10⁵·(e/T²)

where:
- p: pressure [Pa]
- T: temperature [K]
- e: water vapor pressure [Pa]

## Bending Angle
Ray bending through atmosphere is computed via Abel transform:

    α(a) = -2a ∫[a→∞] (dln(n)/dx) / √(x² - a²) dx

where:
- a: impact parameter (distance of ray from Earth center)
- n: refractive index = 1 + N×10⁻⁶
- x: radial distance

# Forward and Adjoint Operators

Forward: (T, q, p) profiles → (N, α) observations
Adjoint: ∂J/∂(N,α) → ∂J/∂(T,q,p)

The adjoint is crucial for data assimilation to propagate observation
information back to model variables.
"""
module GPSRO

using LinearAlgebra
using ..Interpolation

# Import interpolation functions
import ..Interpolation: vertical_interpolate, vertical_interpolate_adjoint!
import ..Interpolation: pressure_to_sigma

export GPSROOperator
export refractivity_forward, refractivity_adjoint
export bending_angle_forward, bending_angle_adjoint
export compute_refractivity, compute_refractivity_gradient
export abel_transform, abel_transform_adjoint

# Physical constants
const Rd = 287.0        # Gas constant for dry air [J/(kg·K)]
const Rv = 461.5        # Gas constant for water vapor [J/(kg·K)]
const k1 = 77.6         # Refractivity constant [K/Pa]
const k2 = 3.73e5       # Refractivity constant [K²/Pa]
const earth_radius = 6.371e6  # Mean Earth radius [m]

"""
    GPSROOperator{T<:AbstractFloat}

GPS Radio Occultation observation operator configuration.

# Fields
- `observable::Symbol`: :refractivity or :bending_angle
- `earth_radius::T`: Earth radius for geometry [m]
- `use_abel_transform::Bool`: Use full Abel transform (true) or simplified (false)
"""
struct GPSROOperator{T<:AbstractFloat}
    observable::Symbol
    earth_radius::T
    use_abel_transform::Bool

    function GPSROOperator{T}(;
                             observable::Symbol=:bending_angle,
                             earth_radius::T=T(6.371e6),
                             use_abel_transform::Bool=true) where T
        @assert observable in [:refractivity, :bending_angle] "Invalid observable type"
        new{T}(observable, earth_radius, use_abel_transform)
    end
end

GPSROOperator(; kwargs...) = GPSROOperator{Float64}(; kwargs...)

# Refractivity Operators

"""
    compute_refractivity(t::T, q::T, p::T) -> T

Compute atmospheric refractivity from temperature, humidity, and pressure.

N = k₁·(p/T) + k₂·(e/T²)

where e is water vapor pressure.
"""
function compute_refractivity(t::T, q::T, p::T) where T<:AbstractFloat
    # Convert specific humidity to water vapor pressure
    # e = q·p / (ε + q·(1-ε)) where ε = Rd/Rv
    epsilon = Rd / Rv
    e = q * p / (epsilon + q * (T(1) - epsilon))

    # Compute refractivity
    n_dry = k1 * p / t
    n_wet = k2 * e / (t * t)

    return n_dry + n_wet
end

"""
    compute_refractivity_gradient(t::T, q::T, p::T) -> Tuple{T,T,T}

Compute gradient of refractivity with respect to T, q, p.

Returns (∂N/∂T, ∂N/∂q, ∂N/∂p)
"""
function compute_refractivity_gradient(t::T, q::T, p::T) where T<:AbstractFloat
    epsilon = Rd / Rv
    e = q * p / (epsilon + q * (T(1) - epsilon))

    # ∂N/∂T
    dN_dT = -k1 * p / (t * t) - T(2) * k2 * e / (t * t * t)

    # ∂N/∂q (via chain rule through e)
    de_dq = p * epsilon / ((epsilon + q * (T(1) - epsilon))^2)
    dN_dq = k2 * de_dq / (t * t)

    # ∂N/∂p
    de_dp = q / (epsilon + q * (T(1) - epsilon))
    dN_dp = k1 / t + k2 * de_dp / (t * t)

    return (dN_dT, dN_dq, dN_dp)
end

"""
    refractivity_forward(t_profile::AbstractVector{T},
                        q_profile::AbstractVector{T},
                        p_profile::AbstractVector{T},
                        z_profile::AbstractVector{T},
                        z_obs::T) -> T

Forward operator for refractivity observations.

Interpolates profiles to observation height and computes refractivity.

# Arguments
- `t_profile`: Temperature profile [K]
- `q_profile`: Specific humidity profile [kg/kg]
- `p_profile`: Pressure profile [Pa]
- `z_profile`: Height profile [m]
- `z_obs`: Observation height [m]

# Returns
- Refractivity at observation height
"""
function refractivity_forward(t_profile::AbstractVector{T},
                             q_profile::AbstractVector{T},
                             p_profile::AbstractVector{T},
                             z_profile::AbstractVector{T},
                             z_obs::T) where T<:AbstractFloat

    # Interpolate profiles to observation height
    t_obs = vertical_interpolate(t_profile, z_profile, z_obs, method=:linear)
    q_obs = vertical_interpolate(q_profile, z_profile, z_obs, method=:linear)
    p_obs = vertical_interpolate(p_profile, z_profile, z_obs, method=:logarithmic)

    # Compute refractivity
    n_obs = compute_refractivity(t_obs, q_obs, p_obs)

    return n_obs
end

"""
    refractivity_adjoint!(t_gradient::AbstractVector{T},
                         q_gradient::AbstractVector{T},
                         p_gradient::AbstractVector{T},
                         residual::T,
                         t_profile::AbstractVector{T},
                         q_profile::AbstractVector{T},
                         p_profile::AbstractVector{T},
                         z_profile::AbstractVector{T},
                         z_obs::T)

Adjoint operator for refractivity observations.

Distributes refractivity gradient back to T, q, p profiles.
"""
function refractivity_adjoint!(t_gradient::AbstractVector{T},
                              q_gradient::AbstractVector{T},
                              p_gradient::AbstractVector{T},
                              residual::T,
                              t_profile::AbstractVector{T},
                              q_profile::AbstractVector{T},
                              p_profile::AbstractVector{T},
                              z_profile::AbstractVector{T},
                              z_obs::T) where T<:AbstractFloat

    # Get interpolated values at observation height
    t_obs = vertical_interpolate(t_profile, z_profile, z_obs, method=:linear)
    q_obs = vertical_interpolate(q_profile, z_profile, z_obs, method=:linear)
    p_obs = vertical_interpolate(p_profile, z_profile, z_obs, method=:logarithmic)

    # Compute refractivity gradients
    dN_dT, dN_dq, dN_dp = compute_refractivity_gradient(t_obs, q_obs, p_obs)

    # Chain rule: multiply by Jacobian
    t_residual = residual * dN_dT
    q_residual = residual * dN_dq
    p_residual = residual * dN_dp

    # Adjoint of vertical interpolation
    vertical_interpolate_adjoint!(t_gradient, t_residual, z_profile, z_obs, method=:linear)
    vertical_interpolate_adjoint!(q_gradient, q_residual, z_profile, z_obs, method=:linear)
    vertical_interpolate_adjoint!(p_gradient, p_residual, z_profile, z_obs, method=:logarithmic)

    return nothing
end

# Bending Angle Operators

"""
    bending_angle_forward(t_profile::AbstractVector{T},
                         q_profile::AbstractVector{T},
                         p_profile::AbstractVector{T},
                         z_profile::AbstractVector{T},
                         impact_param::T;
                         earth_radius::T=T(6.371e6),
                         use_abel::Bool=true) -> T

Forward operator for bending angle observations.

Computes ray bending angle via Abel transform of refractivity profile.

# Arguments
- `t_profile`, `q_profile`, `p_profile`: Atmospheric profiles
- `z_profile`: Height profile [m]
- `impact_param`: Impact parameter (ray closest approach) [m]
- `earth_radius`: Earth radius [m]
- `use_abel`: Use full Abel transform or simplified

# Returns
- Bending angle [radians]
"""
function bending_angle_forward(t_profile::AbstractVector{T},
                               q_profile::AbstractVector{T},
                               p_profile::AbstractVector{T},
                               z_profile::AbstractVector{T},
                               impact_param::T;
                               earth_radius::T=T(6.371e6),
                               use_abel::Bool=true) where T<:AbstractFloat

    nlev = length(t_profile)

    # Compute refractivity profile
    n_profile = zeros(T, nlev)
    for k in 1:nlev
        n_profile[k] = compute_refractivity(t_profile[k], q_profile[k], p_profile[k])
    end

    if use_abel
        # Full Abel transform
        alpha = abel_transform(n_profile, z_profile, impact_param, earth_radius)
    else
        # Simplified: linear relationship
        # α ≈ 10⁻⁶ × N (very crude approximation)
        k_closest = argmin(abs.(z_profile .+ earth_radius .- impact_param))
        alpha = n_profile[k_closest] * T(1e-6)
    end

    return alpha
end

"""
    bending_angle_adjoint!(t_gradient::AbstractVector{T},
                          q_gradient::AbstractVector{T},
                          p_gradient::AbstractVector{T},
                          residual::T,
                          t_profile::AbstractVector{T},
                          q_profile::AbstractVector{T},
                          p_profile::AbstractVector{T},
                          z_profile::AbstractVector{T},
                          impact_param::T;
                          earth_radius::T=T(6.371e6),
                          use_abel::Bool=true)

Adjoint operator for bending angle observations.

Distributes bending angle gradient back to T, q, p profiles via adjoint Abel transform.
"""
function bending_angle_adjoint!(t_gradient::AbstractVector{T},
                               q_gradient::AbstractVector{T},
                               p_gradient::AbstractVector{T},
                               residual::T,
                               t_profile::AbstractVector{T},
                               q_profile::AbstractVector{T},
                               p_profile::AbstractVector{T},
                               z_profile::AbstractVector{T},
                               impact_param::T;
                               earth_radius::T=T(6.371e6),
                               use_abel::Bool=true) where T<:AbstractFloat

    nlev = length(t_profile)

    # Compute refractivity profile
    n_profile = zeros(T, nlev)
    for k in 1:nlev
        n_profile[k] = compute_refractivity(t_profile[k], q_profile[k], p_profile[k])
    end

    # Adjoint of Abel transform
    n_gradient = zeros(T, nlev)

    if use_abel
        n_gradient = abel_transform_adjoint(residual, n_profile, z_profile,
                                           impact_param, earth_radius)
    else
        # Simplified adjoint
        k_closest = argmin(abs.(z_profile .+ earth_radius .- impact_param))
        n_gradient[k_closest] = residual * T(1e-6)
    end

    # Adjoint of refractivity computation
    # Chain rule: ∂J/∂T = ∂J/∂N · ∂N/∂T
    for k in 1:nlev
        dN_dT, dN_dq, dN_dp = compute_refractivity_gradient(t_profile[k],
                                                            q_profile[k],
                                                            p_profile[k])

        t_gradient[k] += n_gradient[k] * dN_dT
        q_gradient[k] += n_gradient[k] * dN_dq
        p_gradient[k] += n_gradient[k] * dN_dp
    end

    return nothing
end

# Abel Transform Implementation

"""
    abel_transform(n_profile::AbstractVector{T},
                  z_profile::AbstractVector{T},
                  impact_param::T,
                  earth_radius::T) -> T

Compute bending angle via Abel transform.

α(a) = -2a ∫[a→∞] (dln(n)/dr) / √(r² - a²) dr

This is simplified using trapezoidal integration.
"""
function abel_transform(n_profile::AbstractVector{T},
                       z_profile::AbstractVector{T},
                       impact_param::T,
                       earth_radius::T) where T<:AbstractFloat

    nlev = length(n_profile)

    # Convert to radial coordinates
    r_profile = z_profile .+ earth_radius

    # Find levels above impact parameter
    indices_above = findall(r -> r >= impact_param, r_profile)

    if isempty(indices_above)
        return T(0)
    end

    # Compute refractive index (n = 1 + N×10⁻⁶)
    refractive_index = T(1) .+ n_profile .* T(1e-6)

    # Numerical integration using trapezoidal rule
    alpha = T(0)

    for i in indices_above[1:end-1]
        r1 = r_profile[i]
        r2 = r_profile[i+1]
        n1 = refractive_index[i]
        n2 = refractive_index[i+1]

        # Skip if below impact parameter
        if r1 < impact_param
            continue
        end

        # Compute integrand at endpoints
        # Integrand: dln(n)/dr / √(r² - a²)
        dln_n_dr1 = (log(n2) - log(n1)) / (r2 - r1)
        dln_n_dr2 = dln_n_dr1  # Constant over interval

        denom1 = sqrt(max(r1^2 - impact_param^2, T(0)))
        denom2 = sqrt(max(r2^2 - impact_param^2, T(0)))

        if denom1 > eps(T) && denom2 > eps(T)
            integrand1 = dln_n_dr1 / denom1
            integrand2 = dln_n_dr2 / denom2

            # Trapezoidal rule
            dr = r2 - r1
            alpha += T(0.5) * (integrand1 + integrand2) * dr
        end
    end

    # Multiply by -2a
    alpha *= -T(2) * impact_param

    return alpha
end

"""
    abel_transform_adjoint(alpha_gradient::T,
                          n_profile::AbstractVector{T},
                          z_profile::AbstractVector{T},
                          impact_param::T,
                          earth_radius::T) -> Vector{T}

Adjoint of Abel transform.

Distributes bending angle gradient back to refractivity profile.

This is the transpose of the Abel transform Jacobian.
"""
function abel_transform_adjoint(alpha_gradient::T,
                               n_profile::AbstractVector{T},
                               z_profile::AbstractVector{T},
                               impact_param::T,
                               earth_radius::T) where T<:AbstractFloat

    nlev = length(n_profile)
    n_gradient = zeros(T, nlev)

    # Convert to radial coordinates
    r_profile = z_profile .+ earth_radius

    # Find levels above impact parameter
    indices_above = findall(r -> r >= impact_param, r_profile)

    if isempty(indices_above)
        return n_gradient
    end

    # Compute refractive index
    refractive_index = T(1) .+ n_profile .* T(1e-6)

    # Adjoint of numerical integration
    # This is a simplified adjoint - distributes gradient uniformly
    # Full adjoint would require careful derivation of integration adjoint

    for i in indices_above[1:end-1]
        r1 = r_profile[i]
        r2 = r_profile[i+1]

        if r1 < impact_param
            continue
        end

        denom1 = sqrt(max(r1^2 - impact_param^2, T(0)))
        denom2 = sqrt(max(r2^2 - impact_param^2, T(0)))

        if denom1 > eps(T) && denom2 > eps(T)
            # Simplified: distribute gradient based on contribution weight
            weight = T(1) / (denom1 + denom2)
            dr = r2 - r1

            # Factor of -2a from Abel transform
            contrib = -T(2) * impact_param * weight * dr * alpha_gradient

            # Distribute to neighboring levels
            n_gradient[i] += contrib * T(0.5)
            n_gradient[i+1] += contrib * T(0.5)
        end
    end

    # Account for N = n×10⁶ scaling
    n_gradient .*= T(1e-6)

    return n_gradient
end

end # module GPSRO
